#include <string>
#include <vector>
#include <unordered_map>
#include <mutex>
#include <functional>
#include <set>

#include "../common/net.hpp"
#include "../common/message.hpp"
#include "../common/uuid.hpp"

namespace JsonMuduoRPC
{
    namespace server
    {
        class ProviderManager
        {
        public:
            using ptr = std::shared_ptr<ProviderManager>;

            struct Provider
            {
                using ptr = std::shared_ptr<Provider>;

                std::mutex mutex;
                std::vector<std::string> methods; // 该提供者可以提供哪些方法
                Address host;                     // 该提供者的主机信息
                BaseConnection::ptr conn;         // 该提供者对应的连接

                Provider(const Address &h, const BaseConnection::ptr &c) : host(h), conn(c) {}

                void appendMethod(const std::string &method)
                {
                    std::unique_lock<std::mutex> lock(mutex);
                    methods.push_back(method);
                }
            };

            // 服务注册时调用
            void addProvider(const BaseConnection::ptr &conn, const Address &host, const std::string &method)
            {
                Provider::ptr provider;

                {
                    std::unique_lock<std::mutex> lock(_mutex);

                    // 1.在_conns_map中找到连接和服务提供者的映射关系，如果没有，就新建
                    auto it = _conns_map.find(conn);
                    if (it == _conns_map.end())
                    {
                        provider = std::make_shared<Provider>(host, conn);
                        _conns_map.insert(std::make_pair(conn, provider));
                    }
                    else
                        provider = it->second;

                    // 2.在_providers_map中添加method方法对应的主机
                    _providers_map[method].insert(provider);
                }

                // 3. 向服务对象中新增一个所能提供的服务名称
                provider->appendMethod(method);
            }

            // 获取一个服务提供者的信息，当一个服务提供者断开连接时，进行服务下线通知时使用
            Provider::ptr getProvider(const BaseConnection::ptr &conn)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                auto it = _conns_map.find(conn);
                if (it == _conns_map.end())
                    return Provider::ptr();
                else
                    return it->second;
            }

            // 当一个服务提供者断开连接的时候，删除它的关联信息
            void deleteProvider(const BaseConnection::ptr &conn)
            {
                std::unique_lock<std::mutex> lock(_mutex);

                // 1. 首先判断该连接对应的是否是一个服务提供者
                auto it = _conns_map.find(conn);
                if (it == _conns_map.end())
                    return;

                // 2. 遍历该连接对应的服务提供者能提供的方法，通过_providers_map删除对应的服务提供者
                Provider::ptr provider = it->second;
                for (auto &method : provider->methods)
                {
                    if (_providers_map.count(method))
                        _providers_map[method].erase(provider);
                }

                // 3. 删除_conns_map映射
                _conns_map.erase(it);
            }

            // 进行服务发现时，返回可以提供method方法的所有主机信息
            std::vector<Address> methodHosts(const std::string &method)
            {
                std::unique_lock<std::mutex> lock(_mutex);

                auto it = _providers_map.find(method);
                if (it == _providers_map.end())
                    return std::vector<Address>();

                std::vector<Address> hosts;
                for (auto &provider : it->second)
                    hosts.push_back(provider->host);

                return hosts;
            }

        private:
            std::mutex _mutex;
            std::unordered_map<std::string, std::set<Provider::ptr>> _providers_map; // 什么方法可以被哪些主机提供
            std::unordered_map<BaseConnection::ptr, Provider::ptr> _conns_map;       // 哪个连接对应哪个提供者
        };

        class DiscoverManager
        {
        public:
            using ptr = std::shared_ptr<DiscoverManager>;

            struct Discover
            {
                using ptr = std::shared_ptr<Discover>;

                std::mutex mutex;
                std::vector<std::string> methods; // 发现过的服务名称
                BaseConnection::ptr conn;         // 发现者关联的连接

                Discover(const BaseConnection::ptr &c) : conn(c)
                {
                }
                void appendMethod(const std::string &method)
                {
                    std::unique_lock<std::mutex> lock(mutex);
                    methods.push_back(method);
                }
            };

            // 当客户端进行服务发现的时候新增发现者，新增服务名称
            Discover::ptr addDiscover(const BaseConnection::ptr &conn, const std::string &method)
            {
                Discover::ptr discover;

                {
                    std::unique_lock<std::mutex> lock(_mutex);

                    // 1.在_conns_map中找到连接和服务发现者的映射关系，如果没有，就新建
                    auto it = _conns_map.find(conn);
                    if (it == _conns_map.end())
                    {
                        discover = std::make_shared<Discover>(conn);
                        _conns_map.insert(std::make_pair(conn, discover));
                    }
                    else
                        discover = it->second;

                    // 2.在_discovers_map中添加method方法对应的主机
                    _discovers_map[method].insert(discover);
                }

                // 3. 向发现对象中新增一个发现了的服务
                discover->appendMethod(method);

                return discover;
            }

            // 发现者客户端断开连接时，找到发现者信息，删除关联数据
            void deleteDiscover(const BaseConnection::ptr &conn)
            {
                std::unique_lock<std::mutex> lock(_mutex);

                // 1. 首先判断该连接对应的是否是一个服务发现者
                auto it = _conns_map.find(conn);
                if (it == _conns_map.end())
                    return;

                // 2. 遍历该连接对应的服务发现者所发现的方法，通过_discovers_map删除对应的服务发现者
                Discover::ptr discover = it->second;
                for (auto &method : discover->methods)
                {
                    if (_discovers_map.count(method))
                        _discovers_map[method].erase(discover);
                }

                // 3. 删除_conns_map映射
                _conns_map.erase(it);
            }

            // 当有一个新的服务提供者上线，进行上线通知
            void onlineNotify(const std::string &method, const Address &host)
            {
                notify(method, host, ServiceOpType::SERVICE_ONLINE);
            }

            // 当有服务提供者下线时，进行下线通知
            void offlineNotify(const std::string &method, const Address &host)
            {
                notify(method, host, ServiceOpType::SERVICE_OFFLINE);
            }

        private:
            void notify(const std::string &method, const Address &host, const ServiceOpType &opType)
            {
                auto it = _discovers_map.find(method);
                // 方法还没有被发现
                if (it == _discovers_map.end())
                    return;

                auto &discovers = it->second;

                // 构建请求
                ServiceRequest::ptr req = MessageFactory::create<ServiceRequest>();
                req->setHost(host);
                req->setId(UUID::getUuidStr());
                req->setMethod(method);
                req->setMType(MType::REQ_SERVICE);
                req->setOpType(opType);

                // 每一个发现了该方法的服务发现者都需要发送请求
                for (auto &discover : discovers)
                    discover->conn->send(req);
            }

        private:
            std::mutex _mutex;
            std::unordered_map<std::string, std::set<Discover::ptr>> _discovers_map; // 哪个方法被哪些发现者发现过
            std::unordered_map<BaseConnection::ptr, Discover::ptr> _conns_map;       // 哪个连接对应哪个发现者
        };

        // 对外
        class ProviderDiscoverManager
        {
        public:
            using ptr = std::shared_ptr<ProviderDiscoverManager>;

            ProviderDiscoverManager() : _provider_manager(std::make_shared<ProviderManager>()),
                                        _discover_manager(std::make_shared<DiscoverManager>()) {}

            // 提供给dispatcher模块的处理服务请求的回调
            // 服务注册/发现
            // 不可能收到服务上线/下线请求，因为服务上线/下线请求只会由服务发现者发起，而服务发现者由该模块管理
            void onServiceRequest(const BaseConnection::ptr &conn, const ServiceRequest::ptr &msg)
            {
                ServiceOpType opType = msg->getOpType();
                if (opType == ServiceOpType::SERVICE_REGISTRY)
                {
                    // 服务注册
                    // 1. 添加服务对象
                    // 2. 进行服务上线通知
                    // 3. 发送响应
                    LOG(INFO, "%s:%d 注册服务 %s", msg->getHost().first.c_str(), msg->getHost().second, msg->getMethod().c_str());

                    _provider_manager->addProvider(conn, msg->getHost(), msg->getMethod());
                    _discover_manager->onlineNotify(msg->getMethod(), msg->getHost());

                    return registryResponse(conn, msg);
                }
                else if (opType == ServiceOpType::SERVICE_DISCOVERY)
                {
                    // 服务发现
                    // 1. 添加服务发现对象
                    // 2. 发送响应
                    LOG(INFO, "客户端要进行 %s 服务发现！", msg->getMethod().c_str());

                    _discover_manager->addDiscover(conn, msg->getMethod());
                    return discoveResponse(conn, msg);
                }
                else
                {
                    LOG(ERROR, "收到服务操作请求，但是操作类型错误！");
                    return errorResponse(conn, msg);
                }
            }

            void onConnectionShutdown(const BaseConnection::ptr &conn)
            {
                auto provider = _provider_manager->getProvider(conn);

                // 如果是一个服务提供者
                if (provider.get() != nullptr)
                {
                    // 服务注册者关闭连接，进行服务下线通知
                    LOG(INFO, "%s:%d 服务下线", provider->host.first.c_str(), provider->host.second);

                    for (auto &method : provider->methods)
                        _discover_manager->offlineNotify(method, provider->host);

                    _provider_manager->deleteProvider(conn);
                }
                _discover_manager->deleteDiscover(conn);
            }

        private:
            // 错误响应
            void errorResponse(const BaseConnection::ptr &conn, const ServiceRequest::ptr &msg)
            {
                ServiceResponse::ptr rsp = MessageFactory::create<ServiceResponse>();
                rsp->setId(msg->getId());
                rsp->setMType(MType::RSP_SERVICE);
                rsp->setOpType(ServiceOpType::SERVICE_UNKNOW);
                rsp->setRCode(RCode::RCODE_INVALID_OPTYPE);

                conn->send(rsp);
            }

            // 服务注册响应
            void registryResponse(const BaseConnection::ptr &conn, const ServiceRequest::ptr &msg)
            {
                ServiceResponse::ptr rsp = MessageFactory::create<ServiceResponse>();
                rsp->setId(msg->getId());
                rsp->setMType(MType::RSP_SERVICE);
                rsp->setOpType(ServiceOpType::SERVICE_REGISTRY);
                rsp->setRCode(RCode::RCODE_OK);

                conn->send(rsp);
            }

            // 服务发现响应
            // 只有服务发现成功rsp才有hosts和method字段
            void discoveResponse(const BaseConnection::ptr &conn, const ServiceRequest::ptr &msg)
            {
                ServiceResponse::ptr rsp = MessageFactory::create<ServiceResponse>();
                std::vector<Address> hosts = _provider_manager->methodHosts(msg->getMethod());

                rsp->setId(msg->getId());
                rsp->setMType(MType::RSP_SERVICE);
                rsp->setOpType(ServiceOpType::SERVICE_DISCOVERY);
                if (hosts.empty())
                {
                    rsp->setRCode(RCode::RCODE_NOT_FOUND_SERVICE);
                    return (void)conn->send(rsp);
                }
                rsp->setMethod(msg->getMethod());
                rsp->setHosts(hosts);
                rsp->setRCode(RCode::RCODE_OK);

                conn->send(rsp);
            }

        private:
            ProviderManager::ptr _provider_manager;
            DiscoverManager::ptr _discover_manager;
        };
    } // namespace server

} // namespace JsonMuduoRPC